#/usr/bin/env python
# -*- coding: utf-8 -*-
#
# Authors: Bruno Alves, n.º 6152; Neuza Figueira, n.º 6036

import psutil

import csv
import sqlite3

from reportlab.platypus import SimpleDocTemplate, Paragraph, Spacer, Table, TableStyle
from reportlab.lib.styles import getSampleStyleSheet
from reportlab.lib.units import inch
from reportlab.lib import colors

import re
import os


class Connection_Search():
	'''Class that allows the detection of active connections on the current machine.''' 

	def get_active_processes(self):
		'''Method to obtain all the running processes on the machine.
		   @return processes -> List of running processes.'''
		processes = psutil.get_pid_list()
		return processes
		

	def get_active_connections(self, process):
		'''Method to obtain the active connections of a given process.
		   @param process -> Process id to scan.
		   @return result -> List of connections.'''
		result = ""

		if psutil.Process(process).get_connections() != []:
			result = str(psutil.Process(process).get_connections())
		else:
			result = "No Active Connections!"

		return result


	def split_data(self, myList):
		'''Method to split and process the data regarding active connections.
		   @param myList -> List of data regarding the active connections.
		   @return newList -> List of splitted data.'''
		newList = []

		for i in range(len(myList)):
			if re.search("fd=", str(myList[i])):
					lista = str(myList[i]).split("fd=")
					a = lista[1].split(' ')[0]
					a = a.replace(",", "").replace("'", "")

			if re.search("family=", str(myList[i])):
				lista = str(myList[i]).split("family=")
				b = lista[1].split(' ')[0]
				b = b.replace(",", "").replace("'", "")

			if re.search("type=", str(myList[i])):
				lista = str(myList[i]).split("type=")
				c = lista[1].split(' ')[0]
				c = c.replace(",", "").replace("'", "")

			if re.search("local_address=", str(myList[i])):
				lista = str(myList[i]).split("local_address=")
				d = lista[1].split(' ')[0]
				e = lista[1].split(',')[1]
				d = d[2:-3]
				e = e.replace(" ", "").replace("'", "").replace(")", "")

			if re.search("remote_address=", str(myList[i])):
				lista = str(myList[i]).split("remote_address=")
				f = lista[1].split(' ')[0]
				f = f.replace(",", "").replace("'", "")

			if re.search("status=", str(myList[i])):
				lista = str(myList[i]).split("status=")
				g = lista[1].split(' ')[0]
				g = g.replace(",", "").replace("'", "")[:-3]

			newList.append([a,b,c,d,e,f,g])

		return newList


	def create_csv_files(self, process_id, connection_list):
		'''Method to create the csv file regarding connections of a process.
		   @param process_id -> ID of the process.
		   @param connection_list -> Connection list.'''

		outfiledir = os.getcwd() + "/CSV"
		if not os.path.exists(outfiledir):
			os.makedirs(outfiledir)

		i = int(process_id)
		filename = open(outfiledir + '/Connection_Scan_Results_' + str(process_id) + '.csv', 'wb')
		spamWriter = csv.writer(filename, delimiter=',')
		spamWriter.writerow(["FD", "FAMILY", "TYPE", "LOCAL ADDRESS", "LOCAL PORT", "REMOTE ADDRESS", "STATUS"])
		myList = []
		newList = []

		if connection_list == "No Active Connections!":
			newList = [('-','-','-','-','-','-','-')]
			pass
		else:
			for i in range(len(connection_list)):
				myList.append(connection_list[i].split(","))	

			print myList

			newList = self.split_data(myList)
	
		for i in range(len(newList)):
			spamWriter.writerow([newList[i][0],newList[i][1],newList[i][2],newList[i][3],newList[i][4],newList[i][5],newList[i][6]])

		filename.close()


	def create_database(self, process_id, connection_list):
		'''Method to create the sqlite3 database regarding the active connections of a process.
		   @param process_id -> ID of the process.
		   @param connection_list -> Connection list.'''
		outfiledir = os.getcwd() + "/DB"

		if not os.path.exists(outfiledir):
			os.makedirs(outfiledir)

		db_name = outfiledir + "/Connection_Scan_Database_" + str(process_id) + ".db"
		connection = sqlite3.connect(str(db_name))

		cursor = connection.cursor()

		cursor.execute('''drop table if exists data''')
		connection.commit()


		cursor.execute('''create table if not exists data (id INTEGER PRIMARY KEY,
                          				       			   fd TEXT,
                          				       			   family TEXT,
                          				       			   type TEXT,
                          				       			   local_address TEXT,
                          				       			   local_port TEXT,
                          				       			   remote_address TEXT,
                          				       			   status TEXT)''')


		myList = []
		newList = []
		i = 1

		if connection_list == "No Active Connections!":
			newList = [('-','-','-','-','-','-','-')]
			pass
		else:
			for i in range(len(connection_list)):
				myList.append(connection_list[i].split(","))	

			print myList

			newList = self.split_data(myList)

		for i in range(len(newList)):
			cmd = "INSERT INTO data VALUES ('{0}','{1}','{2}','{3}','{4}','{5}','{6}','{7}');".format(i,newList[i][0],newList[i][1],newList[i][2],newList[i][3],newList[i][4],newList[i][5],newList[i][6])

			cursor.execute(cmd)

			i = i + 1
			pass

		#cursor.execute("SELECT * FROM data")
		#rows = cursor.fetchall()
	#	for row in rows:
		#	print row


	def create_pdf(self, process_id, connection_list):
		'''Method to create the pdf file regarding the active connections of a process.
		    @param process_id -> ID of the process.
		    @param connection_list -> Connection list.'''
		styles = getSampleStyleSheet()

		outfilename = "Connection_Scan_" + str(process_id) + ".pdf"
		outfiledir = os.getcwd() + "/PDF"
		outfilepath = os.path.join( outfiledir, outfilename )

		if not os.path.exists(outfiledir):
			os.makedirs(outfiledir)

		headerStyle = styles["Heading1"]

		doc = SimpleDocTemplate(outfilepath)
		Story = []
		style = styles["Normal"]
		title = "Connection Scan of Process " + str(process_id) + ": " + psutil.Process(int(process_id)).name
		p = Paragraph(title, styles["Heading1"])

		Story.append(p)
		Story.append(Spacer(1,0.5*inch))

		data = [("FD", "FAMILY", "TYPE", "LOCAL ADDR", "LOCAL PORT", "REMOTE ADDR", "STATUS")]

		myList = []
		newList = []
		i = 1

		if connection_list == "No Active Connections!":
			newList = [('-','-','-','-','-','-','-')]
			pass
		else:
			for i in range(len(connection_list)):
				myList.append(connection_list[i].split(","))	

			print myList

			newList = self.split_data(myList)

		
		for i in range(len(newList)):
			data.append([newList[i][0],newList[i][1],newList[i][2],newList[i][3],newList[i][4],newList[i][5],newList[i][6]])

		table = Table(data, colWidths=75, rowHeights=30)

		table.setStyle(TableStyle([('INNERGRID', (0,0), (-1, -1), 0.25, colors.black), 
                       ('BOX', (0, 0), (-1, -1), 0.25, colors.black), 
                       ('ALIGN',(0,0),(-1,-1),'CENTER'),
                       ('VALIGN',(0,0),(-1,-1),'MIDDLE')
                     ]))


		Story.append(table)
		doc.build(Story) 
